# infer_operations.py：
#推理菜单相关操作
import tkinter as tk
from tkinter import filedialog, messagebox
import os
import json
import cv2
from PIL import Image
from ui_utils import display_image_on_canvas
from yolopredict import PredictionParamsGUI
from tkinter import Toplevel
from predict_operations import YOLOPredictor
from shared import SharedData

# shared_data = SharedData()

def on_pred_file_select(shared_data, app, event):
    """预测结果文件列表选择事件"""
    pred_listbox = shared_data.ui_elements["pred_listbox"]
    selection = pred_listbox.curselection()
    if not selection:
        return
    index = selection[0]

    # 检查是否使用内存图像模式
    if hasattr(shared_data, 'pred_results') and shared_data.pred_results:
        # 使用内存中的 PIL 图像
        if 0 <= index < len(shared_data.pred_results):
            _, pil_image = shared_data.pred_results[index]
            display_image_on_canvas(pil_image, shared_data.ui_elements["result_canvas"])
            app.show_status(f"显示预测结果: {shared_data.pred_results[index][0]}")
        else:
            app.show_status("错误：选中项索引越界")
        return

    # 兼容旧的文件路径模式（可选）
    if hasattr(shared_data, 'pred_file_paths') and shared_data.pred_file_paths:
        if 0 <= index < len(shared_data.pred_file_paths):
            file_path = shared_data.pred_file_paths[index]
            if os.path.exists(file_path):
                try:
                    image = Image.open(file_path)
                    display_image_on_canvas(image, shared_data.ui_elements["result_canvas"])
                    app.show_status(f"显示预测结果: {os.path.basename(file_path)}")
                except Exception as e:
                    app.show_status(f"加载失败: {e}")
            else:
                app.show_status("文件不存在")
                messagebox.showwarning("警告", "预测结果文件已丢失。")
        else:
            app.show_status("索引越界")
    else:
        app.show_status("无预测结果")


def on_pred_key_up(shared_data, app, event):
    """预测结果列表上键事件"""
    pred_listbox = shared_data.ui_elements["pred_listbox"]
    current = pred_listbox.curselection()
    if not current:
        return
    index = current[0]
    if index > 0:
        pred_listbox.selection_clear(index)
        pred_listbox.selection_set(index - 1)
        pred_listbox.see(index - 1)
        on_pred_file_select(shared_data, app, None)


def on_pred_key_down(shared_data, app, event):
    """预测结果列表下键事件"""
    pred_listbox = shared_data.ui_elements["pred_listbox"]
    current = pred_listbox.curselection()
    if not current:
        return
    index = current[0]
    if index < pred_listbox.size() - 1:
        pred_listbox.selection_clear(index)
        pred_listbox.selection_set(index + 1)
        pred_listbox.see(index + 1)
        on_pred_file_select(shared_data, app, None)

def cmd_infer_params(app):
    """
    打开“预测参数”设置窗口
    app: 主窗口实例（AIApplication），用于获取 shared_data 等信息
    """
    # 创建一个子窗口（Toplevel）
    param_window = Toplevel(app)
    param_window.transient(app)  # 设置为临时窗口（始终在主窗口上方）
    param_window.grab_set()      # 模态：焦点锁定在此窗口，主窗口不可操作

    # 初始化参数设置界面
    gui = PredictionParamsGUI(param_window, app.shared_data, app)

    # 注入回调
    app.shared_data.on_start_inference = lambda: cmd_infer_start(app.shared_data, app)

    # 可选：设置关闭回调
    def on_close():
        param_window.destroy()
        app.show_status("预测参数窗口已关闭")

    param_window.protocol("WM_DELETE_WINDOW", on_close)
    
    app.show_status("调整推理参数")

def cmd_infer_start(shared_data, app):
    """
    推理菜单 -> 开始：对当前在“源文件展示”区域显示的文件（图片、视频或摄像头）执行YOLO预测。
    """
    # 1. 停止任何正在播放的视频或摄像头，避免资源冲突
    if shared_data.video_playing and shared_data.cap:
        shared_data.video_playing = False
        shared_data.cap.release()
        shared_data.cap = None
        app.show_status("已停止视频/摄像头播放")

    # 2. 获取当前源类型
    source_type = shared_data.current_source
    if source_type is None:
        app.show_status("错误：没有有效的源文件可供预测")
        messagebox.showwarning("警告", "请先通过“文件”菜单打开一个图片、视频或摄像头。")
        return

    # 3. 加载模型路径和配置参数
    model_path = os.path.abspath(
        os.path.join(
            os.path.dirname(__file__), "..", "results", "result111", "train", "weights", "best.pt"
        )
    )
    if not os.path.exists(model_path):
        app.show_status(f"错误：模型文件不存在 -> {model_path}")
        messagebox.showerror("错误", f"未找到模型文件：\n{model_path}\n请确认路径正确。")
        return

    # 4. 初始化 YOLOPredictor
    try:
        app.show_status("正在初始化 YOLO 模型...")
        predictor = YOLOPredictor(model_path)
        app.predictor = predictor  # 存储 predictor 实例

        # 5. 根据源类型选择推理方式
        if source_type == "picture":
            source_path = shared_data.image_paths[0]
            app.show_status("正在对图片进行预测...")
            predictor.predict_img([source_path])

            if not predictor.results:
                app.show_status("预测失败：未获得任何结果")
                messagebox.showwarning("警告", "模型预测未返回结果。")
                return

            result = predictor.results[0]
            plotted_img = result.plot()  # BGR
            pil_img = to_pil_image(plotted_img, mode="BGR")  # ✅ 统一转PIL
            result_name = f"{os.path.splitext(os.path.basename(source_path))[0]}_pred"

            shared_data.pred_results = [(result_name, pil_img)]
            pred_listbox = shared_data.ui_elements["pred_listbox"]
            pred_listbox.delete(0, tk.END)
            pred_listbox.insert(tk.END, result_name)
            result_canvas = shared_data.ui_elements["result_canvas"]
            display_image_on_canvas(pil_img, result_canvas)
            app.show_status(f"图片预测完成: {result_name}")


        elif source_type == "video":
            source_path = shared_data.video_paths[0]
            app.show_status("正在对视频进行逐帧预测...")
            frame_results = []

            predictor.predict_video(
                video_path=source_path,
                save_output=False,
                max_frames=None
            )

            for result in predictor.results:
                plotted_img = result.plot()  # BGR
                pil_img = to_pil_image(plotted_img, mode="BGR")  # ✅ 统一
                frame_results.append(pil_img)

            if not frame_results:
                app.show_status("视频预测失败：未获得任何结果")
                messagebox.showwarning("警告", "模型预测未返回结果。")
                return

            # 改为：
            result_name = f"{os.path.splitext(os.path.basename(source_path))[0]}_pred"
            shared_data.pred_results = [(result_name, frame_results[0])]
            pred_listbox = shared_data.ui_elements["pred_listbox"]
            pred_listbox.delete(0, tk.END)
            pred_listbox.insert(tk.END, result_name)
            result_canvas = shared_data.ui_elements["result_canvas"]
            display_image_on_canvas(frame_results[0], result_canvas)
            shared_data.current_frame_index = 0
            shared_data.frame_results = frame_results
            shared_data.video_playing = True
            app.after(33, lambda: update_result_frame(shared_data, app))
            app.show_status(f"视频预测完成，共处理 {len(frame_results)} 帧，正在播放结果...")

        # elif source_type == "camera":
        #     app.show_status("正在对摄像头进行实时预测...")
        #     shared_data.cap = cv2.VideoCapture(0)
        #     if not shared_data.cap.isOpened():
        #         app.show_status("无法打开摄像头")
        #         messagebox.showerror("错误", "无法打开摄像头。")
        #         return
        #     shared_data.current_frame_index = 0
        #     shared_data.video_playing = True
        #     app.after(33, lambda: infer_camera_frame(shared_data, app, predictor))
        elif source_type == "camera":
            app.show_status("正在对摄像头进行实时预测...")

            # ✅ 使用 YOLOPredictor 的 predict_camera 方法
            predictor.predict_camera(
                camera_index=0,
                save_output=False,
                max_duration=None  # 可设置为 60 秒测试
            )

            # ✅ 此时 all_results_data 已被填充
            shared_data.pred_results = []
            for i, result in enumerate(predictor.results):
                # plotted_img = result.plot()
                # rgb_img = cv2.cvtColor(plotted_img, cv2.COLOR_BGR2RGB)
                # pil_img = Image.fromarray(rgb_img)
                plotted_img = result.plot()  # BGR
                pil_img = to_pil_image(plotted_img, mode="BGR")  # ✅ 统一
                shared_data.pred_results.append((f"摄像头帧{i+1}", pil_img))

            # 显示第一帧
            if shared_data.pred_results:
                result_canvas = shared_data.ui_elements["result_canvas"]
                display_image_on_canvas(shared_data.pred_results[0][1], result_canvas)
                pred_listbox = shared_data.ui_elements["pred_listbox"]
                pred_listbox.delete(0, tk.END)
                pred_listbox.insert(tk.END, "摄像头实时画面")

            app.show_status("✅ 摄像头预测完成")

        else:
            app.show_status(f"错误：不支持的源类型 ({source_type})")
            messagebox.showwarning("警告", f"源类型不受支持。")

    except Exception as e:
        app.show_status(f"推理失败: {str(e)}")
        messagebox.showerror("错误", f"推理过程中出错：\n{str(e)}")

    # ✅ 推理成功后，生成并显示详细报告
    try:
        # 获取 predictor 实例
        predictor = app.predictor
        if not hasattr(predictor, 'generate_report'):
            app.show_status("模型不支持生成报告")
            return

        # 生成报告文本
        report_text = predictor.generate_report()

        # 显示报告窗口（使用自定义的可滚动窗口）
        show_report_window(app, "YOLO 模型预测详细信息", report_text)        
        app.show_status("✅ 推理完成，详细报告已生成")

    except Exception as e:
        app.show_status(f"生成报告失败: {str(e)}")
        messagebox.showerror("错误", f"生成报告时出错：\n{str(e)}")

def _show_inference_report(app, predictor):
    """内部函数：生成并显示推理报告"""
    try:
        report_text = predictor.generate_report()
        show_report_window(app, "YOLO 预测详细信息", report_text)
        app.show_status("✅ 详细预测信息已生成")
    except Exception as e:
        app.show_status(f"生成报告失败: {str(e)}")
        messagebox.showerror("错误", f"生成详细信息时出错：\n{str(e)}")

# ✅ 工具函数：OpenCV(BGR/RGB) -> PIL
def to_pil_image(img, mode="BGR"):
    """
    将OpenCV图像转换为PIL.Image
    mode: "BGR" 表示输入是OpenCV默认的BGR，"RGB" 表示输入已经是RGB
    """
    if mode == "BGR":
        rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        rgb_img = img
    return Image.fromarray(rgb_img)

def update_result_frame(shared_data, app):
    """定时更新预测结果画布上的帧"""
    if not shared_data.video_playing or not hasattr(shared_data, 'frame_results') or len(shared_data.frame_results) == 0:
        return

    # ✅ 添加暂停检查
    if shared_data.infer_paused:
        return  # 跳过本次处理

    # ✅ 从 shared_data.ui_elements 获取 result_canvas
    canvas = shared_data.ui_elements["result_canvas"]
    current_index = shared_data.current_frame_index
    total_frames = len(shared_data.frame_results)

    # 显示当前帧
    pil_img = shared_data.frame_results[current_index]
    display_image_on_canvas(pil_img, canvas)

    # 更新索引
    shared_data.current_frame_index += 1
    if shared_data.current_frame_index >= total_frames:
        shared_data.current_frame_index = 0  # 循环播放

    # 下一帧
    app.after(33, lambda: update_result_frame(shared_data, app))

def infer_camera_frame(shared_data, app, predictor):
    """定时从摄像头读取一帧，进行推理，并在 result_canvas 上显示结果"""
    if not shared_data.video_playing or not hasattr(shared_data, 'cap') or not shared_data.cap.isOpened():
        return

    if shared_data.infer_paused:
        return  # 暂停时跳过

    # 1. 从摄像头读取一帧（BGR）
    ret, frame = shared_data.cap.read()
    if not ret:
        app.show_status("无法从摄像头读取帧")
        return

    # 2. YOLO 模型推理（直接用BGR输入）
    result = predictor.model(frame, verbose=False)[0]
    plotted_img = result.plot()  # BGR

    # 3. 转为 PIL.Image (RGB)
    pil_img = to_pil_image(plotted_img, mode="BGR")

    # 4. 在画布上显示
    result_canvas = shared_data.ui_elements["result_canvas"]
    display_image_on_canvas(pil_img, result_canvas)

    # 5. 下一帧
    app.after(33, lambda: infer_camera_frame(shared_data, app, predictor))

    
def cmd_infer_pause(app):
    """
    推理菜单 -> 暂停：暂停当前的推理过程。
    """
    # 获取 shared_data
    shared_data = app.shared_data

    if not hasattr(shared_data, 'video_playing') or not shared_data.video_playing:
        app.show_status("当前没有正在进行的推理")
        messagebox.showinfo("提示", "当前没有推理任务在运行。")
        return

    # 切换暂停状态
    if shared_data.infer_paused:
        # 已经暂停，现在要恢复
        shared_data.infer_paused = False
        app.show_status("已恢复推理")

        # 如果是摄像头，重新启动定时器
        if shared_data.current_source == "camera" and hasattr(shared_data, 'cap') and shared_data.cap.isOpened():
            app.after(33, lambda: infer_camera_frame(shared_data, app, app.predictor))
    else:
        # 正在运行，现在要暂停
        shared_data.infer_paused = True
        app.show_status("已暂停推理")


def cmd_infer_stop(app):
    """
    推理菜单 -> 停止：完全停止推理并清理资源。
    """
    shared_data = app.shared_data

    # 停止任何正在进行的推理
    if hasattr(shared_data, 'video_playing') and shared_data.video_playing:
        shared_data.video_playing = False
        shared_data.infer_paused = False  # 确保暂停标志被清除

        # 停止摄像头
        if hasattr(shared_data, 'cap') and shared_data.cap is not None:
            shared_data.cap.release()
            shared_data.cap = None

        # 清空 result_canvas
        result_canvas = shared_data.ui_elements["result_canvas"]
        result_canvas.delete("all")
        result_canvas.create_text(result_canvas.winfo_width() // 2,
                                  result_canvas.winfo_height() // 2,
                                  text="推理已停止\n请重新开始",
                                  fill="gray",
                                  font=("SimHei", 12) if "SimHei" in tk.font.families() else ("Arial", 12),
                                  tags="placeholder")

        # 清空 pred_results 和 frame_results
        if hasattr(shared_data, 'pred_results'):
            shared_data.pred_results = []
        if hasattr(shared_data, 'frame_results'):
            shared_data.frame_results = []

        app.show_status("推理已停止")

def show_report_window(parent, title, text):
    """
    创建一个显示长文本报告的窗口
    """
    window = Toplevel(parent)
    window.title(title)
    window.geometry("850x600")
    window.transient(parent)
    window.grab_set()

    # 创建文本框和滚动条
    text_frame = tk.Frame(window)
    text_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

    text_widget = tk.Text(text_frame, wrap=tk.WORD, font=("Courier", 10), bg="#f8f8f8", fg="#333")
    scrollbar = tk.Scrollbar(text_frame, orient=tk.VERTICAL, command=text_widget.yview)
    text_widget.configure(yscrollcommand=scrollbar.set)

    text_widget.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
    scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

    # 插入文本
    text_widget.insert(tk.END, text)
    text_widget.config(state=tk.DISABLED)  # 只读

    # 添加关闭按钮
    button_frame = tk.Frame(window)
    button_frame.pack(pady=5)
    tk.Button(button_frame, text="关闭", command=window.destroy, width=10).pack()

    # 窗口居中
    window.update_idletasks()
    x = parent.winfo_x() + (parent.winfo_width() // 2) - (window.winfo_width() // 2)
    y = parent.winfo_y() + (parent.winfo_height() // 2) - (window.winfo_height() // 2)
    window.geometry(f"+{x}+{y}")